#---------------------------------------------------------------------------#
#       end-to-end网络：输入层 + encoder_net + decoder_net + 输出层 + 生成模型
#---------------------------------------------------------------------------#
import tensorflow as tf 
import tensorflow.keras as keras

'''
    # 网络建模注意
        1.与原文一致，采用了4次下采样过程
        2.原文下采样和上采样过程，不是严格的尺寸倍增过程
            网络输入尺寸为572，网络输出尺寸为388
            bubbling和本处，使用了严格的尺寸倍增过程 i.e. 网络输入的尺寸，在网络输出时被严格还原
        3.encoder使用了迁移学习的mobilenet模型，没有严格还原unet网络结构
            unet Block_1为尺寸不变的两次卷积，这一部分无法用迁移学习模型还原
            导致了decoder最终输出的尺寸为原输入尺寸的一半，即只能做到3次上采样
            因为，photo没有卷积过程，直接融合到上采样结果，应该弊大于利
        4.unet使用的上采样算法未知，这里使用的是近邻插值
'''
def f(photo, num_classes):
    ###----- 主干特征提取网络：encoder_block
    m_net = keras.applications.mobilenet.MobileNet(include_top=False, 
                                                   weights='imagenet', 
                                                   input_shape=photo.shape[1:4])
    m_net_BlockendLayerIndex = [10,23,36,73,86] # blockIdx,multiple,layerIdx: 0,2×,10  1,4×,23  2,8×,36  3,16×,73  4,32×,86
    
    encoder_block_1 = keras.models.Sequential(m_net.layers[0:10])
    encoder_block_1.trainable = False
    o1 = encoder_block_1(photo)

    encoder_block_2 = keras.models.Sequential(m_net.layers[10:23])
    encoder_block_2.trainable = False
    encoder_block_2.build(input_shape=o1.shape)
    o2 = encoder_block_2(o1)
    
    encoder_block_3 = keras.models.Sequential(m_net.layers[23:36])
    encoder_block_3.trainable = False
    encoder_block_3.build(input_shape=o2.shape)
    o3 = encoder_block_3(o2)
    
    encoder_block_4 = keras.models.Sequential(m_net.layers[36:73])
    encoder_block_4.trainable = False
    encoder_block_4.build(input_shape=o3.shape)
    o4 = encoder_block_4(o3)

    ###----- 加强特征提取网络：decoder_block
    channels = [64, 128, 256, 512]

    o = keras.layers.UpSampling2D(size=(2,2))(o4)
    o = tf.concat([o3, o], axis=-1)
    o = keras.layers.Conv2D(channels[3], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)
    o = keras.layers.Conv2D(channels[3], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)

    o = keras.layers.UpSampling2D(size=(2,2))(o)
    o = tf.concat([o2, o], axis=-1)
    o = keras.layers.Conv2D(channels[2], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)
    o = keras.layers.Conv2D(channels[2], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)

    o = keras.layers.UpSampling2D(size=(2,2))(o)
    o = tf.concat([o1, o], axis=-1)
    o = keras.layers.Conv2D(channels[1], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)
    o = keras.layers.Conv2D(channels[1], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)

    # o = keras.layers.UpSampling2D(size=(2,2))(o)
    # o = tf.concat([photo, o], axis=-1)
    # o = keras.layers.Conv2D(channels[0], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)
    # o = keras.layers.Conv2D(channels[0], (3,3), padding='same', activation='relu', kernel_initializer='he_normal')(o)

    ###----- 预测网络
    o = keras.layers.Conv2D(num_classes+1, (1,1), padding='same')(o)

    return o

def get_net_model(net_input_shape, num_classes): #, num_upSample=3, encoder_level=3):
    # 1.输入层
    i_put = keras.layers.Input(net_input_shape)

    # 2.主干特征提取网络 + 加强特征提取网络
    decoded_map = f(i_put, num_classes)

    # 4.输出层
    o_put = tf.image.resize(decoded_map, net_input_shape[0:2]) # 还原到网络输入尺寸！
    o_put = keras.layers.Softmax()(o_put)

    # 5.模型实例
    net_model = keras.models.Model(i_put, o_put)

    return net_model

if __name__ == '__main__':
    print('TST MSG: 测试Unet网络建模是否正确')
    net_model = get_net_model((480,480,3), 2)
    net_model.summary()